# -*- coding: utf-8 -*-
"""
Created on Thu Sep  5 19:36:07 2024

@author: LENOVO
"""

import matplotlib.pyplot as plt
import numpy as np
from sympy import *
from scipy.optimize import root, fsolve
import pandas as pd

#常数的确定
a=16*55/100  #m
b=55/(2*np.pi)/100 #m
vh=1 #m/s
L0=(341-27.5*2)/100
Lb=(220-27.5*2)/100
# theta=np.linspace(0,32*np.pi,30*180)
r=lambda theta:((a-b*(theta)))

# ax=plt.subplot(111, polar=True)
# ax.set_theta_direction(-1)
# plt.plot(theta,r(theta),lw=1,c='r', label='原始数据点')
# plt.legend()
# plt.title('非线性最小二乘拟合')
# plt.grid()  
# plt.show()

#求解theta与t的关系
# var('theta')
# t=integrate(sqrt((a-b*theta)**2+b**2),theta)
# print(t)

# print("FENGE")
# theta=symbols('theta')
# t=Function("t")
# eq1=diff(t(theta),theta,1)-sqrt(((a-b*theta)**2+b**2))
# s1=dsolve(eq1,t(theta))
# print(s1)


s=lambda theta:a*theta-b*theta**2/2
theta=lambda t:(a-np.sqrt(a**2-2*b*vh*t))/b
t=np.linspace(1,300,300)
# plt.plot(t,s(theta(t)),lw=1,c='r', label='原始数据点')
# plt.plot(t,theta(t),lw=1,c='r', label='原始数据点')
# plt.plot(t,r(theta(t)),lw=1,c='r', label='原始数据点')
# plt.legend()
# plt.title('非线性最小二乘拟合')
# plt.grid()  
# plt.show()



#循环遍历
# for i in range(300,-1,-1):    #300s到0s循环遍历
#第i秒位置函数

def position(i):
    Data=[]


    THETA=np.ones((301,1))
    position=np.ones((301,2))
    # y=np.ones((301,1))
    
    THETA[i]=theta(i)
    position[i,0]=r(THETA[i])*np.cos((THETA[i]))
    position[i,1]=-r(THETA[i])*np.sin((THETA[i]))

    Data.append(theta(i))
    f=lambda thetai:(r(THETA[i]))**2+(r(thetai))**2-L0**2-2*(r(THETA[i]))*(r(thetai))*np.cos(THETA[i]-thetai)
    thetai=fsolve(f,THETA[i]-0.5)
    
    # print("theta1:",thetai)
    THETA[i-1]=thetai
    position[i-1,0]=r(THETA[i-1])*np.cos((THETA[i-1]))
    position[i-1,1]=-r(THETA[i-1])*np.sin((THETA[i-1]))
    Data.append(thetai)
    
    for j in range(0,221,1):
        
        f=lambda thetai:(r(THETA[i-1-j]))**2+(r(thetai))**2-Lb**2-2*(r(THETA[i-1-j]))*(r(thetai))*np.cos(THETA[i-1-j]-thetai)
        thetai=fsolve(f,THETA[i-1-j]-0.5)
        THETA[i-2-j]=thetai
        position[i-2-j,0]=r(THETA[i-2-j])*np.cos((THETA[i-2-j]))
        position[i-2-j,1]=-r(THETA[i-2-j])*np.sin((THETA[i-2-j]))
        if THETA[i-2-j]>=0:
            Data.append(thetai)
        else:
            print("龙身进入了:",j+1)
            break
    # print(Data)

    # print(position)
    X=position[:,0]
    Y=position[:,1]
    THETA=[x for x in THETA if(x!=1)]
    position=[(x,y) for (x,y) in position if (x!=1 and y!=1)]
    X=[x for x in X if(x!=1)]
    Y=[x for x in Y if(x!=1)]
    THETA=np.array(THETA)
    X=np.array(X)
    Y=np.array(Y)
    THETA=THETA.reshape(-1,1)
    X=X.reshape(-1,1)
    Y=Y.reshape(-1,1)
    # position=position.reshape(-1,1)
    A=np.column_stack((X, Y, THETA))
    # print(A)
    return A
    
WZ=position(442)
print(WZ)    
# N=[]
# for k in range (0,301,1):
#     N.append(position(k))
# print(N)

# #导出到excel
# for k in range (0,301,1):

    
# df.to_excel('output.xlsx',index=False)

# df=pd.DataFrame(position(300))
# df.to_excel('300position.xlsx',index=False)
    
    
    
# f=lambda theta1:(r(theta(300)))**2+(r(theta1))**2-L0**2-2*(r(theta(300)))*(r(theta1))*np.cos(theta(300)-theta1)
# theta1=fsolve(f,43)
# print("theta1:",theta1[0])
# THETA=[]
# THETA.append(theta1)
# print(theta(300))
# print(s(theta(300)))


def Slope(ta):
    k=(-a*np.cos(ta)+b*np.sin(ta)+b*ta*np.cos(ta))/-a*np.sin(ta)-b*np.cos(ta)+b*ta*np.sin(ta)
    return k

#速度的求解 #第i秒时候所有的速度
def Velocity(i):
    v=0
    A=np.array([1,Slope(theta(i))])/np.sqrt(1**2+Slope(theta(i))**2)
    return A
    
    
    




